# -*- coding:utf-8 -*-

from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex
import json


def repair_to_16bit(text):
    if len(text.encode('utf-8')) % 16:
        add = 16 - (len(text.encode('utf-8')) % 16)
    else:
        add = 0
    text = text + ('\0' * add)
    return text.encode('utf-8')


def encrypt(data: str or dict):
    if isinstance(data, dict):
        text = json.dumps(data, ensure_ascii=False)
    else:
        text = data
    key = '5tyjkeds06.hg*cv@xsq!在'.encode('utf-8')
    mode = AES.MODE_CBC
    iv = b'sdftgbklfxdgtu9h'
    text = repair_to_16bit(text)
    cryptos = AES.new(key, mode, iv)
    cipher_text = cryptos.encrypt(text)
    # 因为AES加密后的字符串不一定是ascii字符集的，输出保存可能存在问题，所以这里转为16进制字符串
    return b2a_hex(cipher_text).decode()


passwd_enc = encrypt
token_enc = encrypt


def decrypt(text):
    key = '5tyjkeds06.hg*cv@xsq!在'.encode('utf-8')
    iv = b'sdftgbklfxdgtu9h'
    mode = AES.MODE_CBC
    cryptos = AES.new(key, mode, iv)
    plain_text = cryptos.decrypt(a2b_hex(text))
    return bytes.decode(plain_text).rstrip('\0')


passwd_dec = decrypt
token_dec = decrypt


if __name__ == '__main__':
    import datetime
    user_token = json.dumps({
        'i': 2,
        'n': 'atp',
        'lt': 0,
        'et': (datetime.datetime.now() + datetime.timedelta(days=36500)).timestamp()
    })
    e = encrypt(user_token)  # 加密
    d = decrypt(e)  # 解密
    print("加密:", e)
    print("解密:", d)



